import math
import random
from collections import Counter

random.seed(42)


def calc_euclidean_dist(s1, s2):
    # 统一校验样本格式与特征类型
    for s, name in [(s1, "样本1"), (s2, "样本2")]:
        if not isinstance(s, list):
            raise TypeError(f"{name}需为列表类型")
        if not all(isinstance(val, (int, float)) for val in s):
            raise TypeError(f"{name}特征值需为整数/浮点数")
    # 校验维度一致性
    if len(s1) != len(s2):
        raise ValueError(f"样本维度不匹配:样本1({len(s1)}维) != 样本2({len(s2)}维)")
    # 计算欧氏距离
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(s1, s2)))


def knn_predict(X_train, y_train, x_test, k=5, task="classification"):
    # 核心输入校验（整合重复逻辑，聚焦关键校验点）
    if not (isinstance(X_train, list) and isinstance(y_train, list) and X_train and y_train):
        raise ValueError("X_train/y_train需为非空列表,且一一对应")
    if len(X_train) != len(y_train):
        raise ValueError(f"X_train({len(X_train)}个样本)与y_train({len(y_train)}个标签)数量不匹配")
    if not (isinstance(k, int) and 0 < k <= len(X_train)):
        raise ValueError(f"k需为满足0 < k ≤ {len(X_train)}的整数")
    if task not in ["classification", "regression"]:
        raise ValueError("task仅支持'classification'(分类), 'regression'(回归)")
    if task == "regression" and not all(isinstance(lab, (int, float)) for lab in y_train):
        raise TypeError("回归任务y_train标签需为整数/浮点数")
    calc_euclidean_dist(X_train[0], x_test)  # 校验测试样本维度

    # 计算距离+筛选top-k近邻
    dist_labels = [(calc_euclidean_dist(x, x_test), lab) for x, lab in zip(X_train, y_train)]
    dist_labels.sort(key=lambda x: (x[0], random.random()))  # 按距离排序，同距随机
    top_k_labels = [lab for _, lab in dist_labels[:k]]

    # 按任务输出结果
    if task == "classification":
        cnt = Counter(top_k_labels)
        max_cnt = max(cnt.values())
        return random.choice([lab for lab, c in cnt.items() if c == max_cnt])
    else:
        return round(sum(top_k_labels) / k, 4)


def knn_evaluate(X_train, y_train, X_test, y_test, k=5, task="classification"):
    # 测试集基础校验
    if not (isinstance(X_test, list) and isinstance(y_test, list) and X_test and y_test):
        raise ValueError("X_test/y_test需为非空列表，且一一对应")
    if len(X_test) != len(y_test):
        raise ValueError(f"X_test（{len(X_test)}个样本）与y_test（{len(y_test)}个标签）数量不匹配")
    if task == "regression" and not all(isinstance(lab, (int, float)) for lab in y_test):
        raise TypeError("回归任务y_test标签需为整数/浮点数")
    calc_euclidean_dist(X_train[0], X_test[0])  # 校验维度一致性

    # 批量预测+计算评估指标
    y_pred = [knn_predict(X_train, y_train, x, k, task) for x in X_test]
    if task == "classification":
        correct = sum(1 for t, p in zip(y_test, y_pred) if t == p)
        acc = correct / len(y_test)
        return {"准确率": round(acc, 4), "正确数/总数": f"{correct}/{len(y_test)}"}
    else:
        mse = sum((t - p) ** 2 for t, p in zip(y_test, y_pred)) / len(y_test)
        return {"均方误差(MSE)": round(mse, 4), "测试样本数": len(y_test)}


if __name__ == "__main__":
    # 鸢尾花分类任务验证
    print("【鸢尾花品种预测】特征：[花瓣长度,花瓣宽度] | 标签：0=山鸢尾,1=变色鸢尾,2=维吉尼亚鸢尾")
    # 数据集
    X_train = [[1.4, 0.2], [1.3, 0.2], [1.5, 0.2], [1.4, 0.3],
               [4.7, 1.4], [4.5, 1.5], [4.9, 1.5], [4.7, 1.6],
               [6.7, 2.2], [6.3, 1.8], [6.5, 2.0], [6.2, 2.3]]
    y_train = [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]
    X_test = [[1.6, 0.2], [4.8, 1.5], [6.4, 2.1], [1.3, 0.3], [5.0, 1.7]]
    y_test = [0, 1, 2, 0, 1]

    # 单样本预测
    single_pred = knn_predict(X_train, y_train, [4.6, 1.4], k=3)
    print(f"样本[4.6,1.4]预测标签：{single_pred}")
    # 测试集评估
    eval_res = knn_evaluate(X_train, y_train, X_test, y_test, k=3)
    print(f"分类评估结果：准确率={eval_res['准确率']}({eval_res['正确数/总数']})")

